package com.agentframework.core.plan;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 计划类
 * @author wangjun
 **/
public class Plan {

    /**
     * 任务目标
     */
    private String goal;

    /**
     * 任务内容
     */
    private String context;

    private List<Task> tasks = new ArrayList<>();

    private Map<String, Task> taskMap = new HashMap<>();

    private String currentTaskId = "";

    private List<Task> topologicalSort(List<Task> tasks) {
        Map<String, Task> taskMap = tasks.stream().collect(Collectors.toMap(Task::getTaskId, task -> task));
        Map<String, Set<String>> dependencies = new HashMap<>();
        tasks.forEach(task -> dependencies.put(task.getTaskId(), new HashSet<>(task.getDependentTaskIds())));

        List<Task> sortedTasks = new ArrayList<>();
        Set<String> visited = new HashSet<>();

        for (Task task : tasks) {
            visit(task.getTaskId(), visited, sortedTasks, dependencies, taskMap);
        }

        return sortedTasks;
    }

    private void visit(String taskId, Set<String> visited, List<Task> sortedTasks, Map<String, Set<String>> dependencies, Map<String, Task> taskMap) {
        if (visited.contains(taskId)) {
            return;
        }
        visited.add(taskId);
        for (String dependentId : dependencies.getOrDefault(taskId, Collections.emptySet())) {
            visit(dependentId, visited, sortedTasks, dependencies, taskMap);
        }
        sortedTasks.add(taskMap.get(taskId));
    }

    public void addTasks(List<Task> newTasks) {
        if (newTasks.isEmpty()) {
            return;
        }

        List<Task> sortedNewTasks = topologicalSort(newTasks);

        if (tasks.isEmpty()) {
            tasks = sortedNewTasks;
        } else {
            int prefixLength = 0;
            for (int i = 0; i < Math.min(tasks.size(), sortedNewTasks.size()); i++) {
                if (!tasks.get(i).getTaskId().equals(sortedNewTasks.get(i).getTaskId()) ||
                        !tasks.get(i).getInstruction().equals(sortedNewTasks.get(i).getInstruction())) {
                    break;
                }
                prefixLength++;
            }

            List<Task> finalTasks = new ArrayList<>(tasks.subList(0, prefixLength));
            finalTasks.addAll(sortedNewTasks.subList(prefixLength, sortedNewTasks.size()));
            tasks = finalTasks;
        }

        updateCurrentTask();
        taskMap = tasks.stream().collect(Collectors.toMap(Task::getTaskId, task -> task));
    }

    public void resetTask(String taskId) {
        if (taskMap.containsKey(taskId)) {
            taskMap.get(taskId).reset();
        }
    }

    public void replaceTask(Task newTask) {
        if (!taskMap.containsKey(newTask.getTaskId())) {
            throw new IllegalArgumentException("Task not found in plan");
        }

        taskMap.put(newTask.getTaskId(), newTask);
        for (int i = 0; i < tasks.size(); i++) {
            if (tasks.get(i).getTaskId().equals(newTask.getTaskId())) {
                tasks.set(i, newTask);
                break;
            }
        }

        tasks.stream()
                .filter(task -> task.getDependentTaskIds().contains(newTask.getTaskId()))
                .forEach(task -> resetTask(task.getTaskId()));
    }

    public void appendTask(Task newTask) {
        if (hasTaskId(newTask.getTaskId())) {
            throw new IllegalArgumentException("Task already in current plan, use replaceTask instead");
        }

        if (!newTask.getDependentTaskIds().stream().allMatch(this::hasTaskId)) {
            throw new IllegalArgumentException("New task has unknown dependencies");
        }

        tasks.add(newTask);
        taskMap.put(newTask.getTaskId(), newTask);
        updateCurrentTask();
    }

    public boolean hasTaskId(String taskId) {
        return taskMap.containsKey(taskId);
    }

    private void updateCurrentTask() {
        currentTaskId = tasks.stream()
                .filter(task -> !task.getFinished())
                .map(Task::getTaskId)
                .findFirst()
                .orElse("");
    }

    public Task getCurrentTask() {
        return taskMap.getOrDefault(currentTaskId, null);
    }

    public void finishCurrentTask() {
        if (currentTaskId != null && !currentTaskId.isEmpty()) {
            Task currentTask = getCurrentTask();
            if (currentTask != null) {
                currentTask.setFinished(true);
                updateCurrentTask();
            }
        }
    }

    public List<Task> getFinishedTasks() {
        return tasks.stream()
                .filter(Task::getFinished)
                .collect(Collectors.toList());
    }

    public void setGoal(String goal) {
        this.goal = goal;
    }

    public String getGoal() {
        return goal;
    }

    public String getContext() {
        return context;
    }

    public List<Task> getTasks() {
        return this.tasks;
    }

}
